import jieba
import jieba.analyse
import matplotlib.pyplot as plt
from wordcloud import ImageColorGenerator, WordCloud
from PIL import Image
import numpy as np


def get_stopwords_set(file):
    with open(file, 'r', encoding="UTF-8") as f:
        stopwords_set = set(f.read().split())
    return stopwords_set

def get_requirement(file):
    with open(file, 'r', encoding="UTF-8") as f:
        data = f.read()
    return data

def segment_requirement(sentence:str, stopwords:set=None):
    jieba.load_userdict("./user_dict.txt")
    # word_list is a generator
    word_list = jieba.cut(sentence=sentence)
    word_frequency = dict()
    if stopwords is None:
        stopwords = set()
    for word in word_list: 
        if word not in stopwords and len(word) > 1:
            word_frequency[word] = word_frequency.get(word, 0) + 1
    return word_frequency

def create_wordcloud(word_frequency):
    img = Image.open("./totoro2.jpg").resize((1600,1644))
    mask = np.array(img)
    img_color = ImageColorGenerator(mask)
    # word_cloud = WordCloud(font_path="STXINGKA.TTF", mask=mask, scale=5, background_color='white')
    word_cloud = WordCloud(font_path="msyh.ttc", mask=mask, scale=5, background_color='white')
    word_cloud.generate_from_frequencies(word_frequency)
    word_cloud.recolor(color_func=img_color)
    return word_cloud

if __name__ == "__main__":
    stopwords_file_path = "./stopwords.txt"
    stopwords = get_stopwords_set(stopwords_file_path)
    other = ["能力", "资格", "工作", "经验", "熟悉", "熟练", "要求", "优先", "具有", "相关",
    "年", "强", "良好", "应用", "基本", "以上", "一下", "以上学历", "一种"]
    for word in other:
        stopwords.add(word)
    for i in range(10):
        stopwords.add(str(i))
    requirement_file_path = "./output/detail_数据挖掘.txt"
    requirement = get_requirement(requirement_file_path)
    word_frequency = segment_requirement(requirement, stopwords=stopwords)
    print("关键词提取>")
    print(jieba.analyse.extract_tags(requirement, withWeight=True))
    word_cloud = create_wordcloud(word_frequency)
    # word_cloud.to_file("./wordcloud.jpg")
    plt.figure(figsize=(8, 8))
    plt.axis('off')
    plt.imshow(word_cloud)
    plt.show()